package metrics;

import java.util.ArrayList;
import java.util.Collection;

import thesis.InfoUnit;
import thesis.DataObject;
import thesis.Summary;

public class JSDivergence {
	public static final double eps = 5e-4;

	public static double getScore(final Summary summary,
			final Collection<DataObject> allData,
			final ArrayList<InfoUnit> units, final double[][] probs) {
		double[] pS = new double[units.size()];
		double[] pO = new double[units.size()];
		double unitSize = 0;
		for (InfoUnit unit : units) {
			unitSize += unit.getWeight();
		}
		double totalUnitsInS = 0, totalUnitsInO = 0;
		for (DataObject t : summary.getMemoryTweets()) {
			for (InfoUnit unit : units) {
				if (probs[t.getInternId()][unit.getInternId()] == 1) {
					totalUnitsInS += unit.getWeight();
				}
			}
		}
		for (DataObject t : allData) {
			for (InfoUnit unit : units) {
				if (probs[t.getInternId()][unit.getInternId()] == 1) {
					totalUnitsInO += unit.getWeight();
				}
			}
		}
		for (int i = 0; i < units.size(); i++) {
			InfoUnit unit = units.get(i);
			// calculate PS
			int hit = 0;
			for (DataObject t : summary.getMemoryTweets()) {
				if (probs[t.getInternId()][unit.getInternId()] == 1) {
					hit += unit.getWeight();
				}
			}
			pS[i] = (hit + eps) / (totalUnitsInS + eps * unitSize);
			// calculate PO
			hit = 0;
			for (DataObject t : allData) {
				if (probs[t.getInternId()][unit.getInternId()] == 1) {
					hit += unit.getWeight();
				}
			}
			pO[i] = (hit + eps) / (totalUnitsInO + eps * unitSize);
		}
		// return 1 - getKL(pO, pS);
		return (getKL(pO, pS) + getKL(pS, pO)) / 2;
	}

	public static double getKL(double[] po, double[] ps) {
		double kl = 0;
		for (int i = 0; i < po.length; i++) {
			kl += po[i] * Math.log(po[i] / ps[i]);
		}
		return kl;
	}
}
